#! /usr/bin/env python

"""
Draws the figures for the single tracking layer in a radial cavity
with a magnetic field.

Used in:
~/Documents/cms/upgrade/occupancy/occupancy01/occupancy01.tex
"""

import os
import tempfile
import shutil
from math import *

## Writes the header of a pspicture LaTeX document.
#
def writeheader(paperheight,paperwidth,xoffset,yoffset,scale):
    out = ""
    out += "\\documentclass{article}\n"
    out += "\\usepackage{pstricks}\n"
    out += "\\usepackage{pst-all}\n"
    out += "\\selectcolormodel{rgb}\n"
    out += "\\usepackage{pspicture}\n"
    out += "\\usepackage[dvips,paperwidth=%10.5fcm,paperheight=%10.5fcm,\n" % (paperwidth,paperheight)
    out += "left=0cm,right=0cm,top=0cm,bottom=0cm,headheight=0pt,headsep=0pt,footskip=0pt,\n"
    out += "]{geometry}\n"
    out += "\\oddsidemargin=-1in\n"
    out += "\\topmargin=-1in\n"
    out += "\\parindent=0pt\n"
    out += "\\begin{document}\n"
    out += "\\pagestyle{empty}\n"
    out += "\\SpecialCoor\n"
    out += "\\psset{xunit=%10.5fcm}\n" % (1./scale)
    out += "\\psset{yunit=%10.5fcm}\n" % (1./scale)
    out += "\\psset{runit=%10.5fcm}\n" % (1./scale)
    out += "\\begin{pspicture}(0,0)(-%10.5f,%10.5f)\n" % (xoffset*scale,(paperheight-yoffset)*scale)
    out += "\\linethickness{0.8pt}\n"
    return out 

## Writes the footer for a pspicture LaTeX file.
#
def writefooter():
    out = ""
    out += "\\end{pspicture}\n"
    out += "\\end{document}\n"
    return out

## Draws the cavity walls (cut-offs) in rho-phi space.
#
def drawcutoffrhophi(rho,edge):
    out = ''
    out += '\\psframe*[linecolor=lightgray](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (-(rho+edge),-(rho+edge),rho+edge,rho+edge)
    out += '\\pscircle[linewidth=0.3pt,fillstyle=solid,fillcolor=white](%10.5f,%10.5f){%10.5f}\n' % (0.0,0.0,rho)
    return out

## Draws the cavity walls (cut-offs) in rho-z space.
#
def drawcutoffrhoz(z,rho,edge):
    out = ''
    out += '\\psframe*[linecolor=lightgray](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (-(z+edge),-(rho+edge),z+edge,rho+edge)
    out += '\\psframe[linewidth=0.3pt,fillstyle=solid,fillcolor=white](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (-z,-rho,z,rho)
    return out

## Draws the tracker in rho-z space
#
def drawtrackerrhoz(z,Deltaz,rho,Deltarho):
    out = ""
    if Deltaz < 0.00001: # if not segmented in z
        if Deltarho < 0.00001:
            Deltarho = 1.5
        out += '\\psframe[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (-z, rho,z,  rho+Deltarho)
        out += '\\psframe[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (-z,-rho,z,-(rho+Deltarho))
    else: # if segmented in z
        pass
        #if Deltarho < 0.00001:
        #    Deltarho = 0.01
        #Nz = int(z/Deltaz)
        #Overz = z - float(Nz)*Deltaz
        #for k in range(Nz):
        #    zmin = float(k)*Deltaz
        #    zmax = zmin+Deltaz
        #    out += '\\psframe[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (zmin,rho,zmax,rho+Deltarho)
        #    out += '\\psframe[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (zmin,-rho,zmax,-(rho+Deltarho))
        #    out += '\\psframe[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (-zmin,rho,-zmax,rho+Deltarho)
        #    out += '\\psframe[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (-zmin,-rho,-zmax,-(rho+Deltarho))
        #    #out += '\\pscircle[linewidth=0.3pt](1.0,1.0){1.0}\n'
    return out

## Draws a straight track in rho-phi space.
#
def drawstraighttrackrhophi(phi,eta,rhomax,zmax):
    etatrans = asinh(zmax/rhomax)
    phi_deg = 180. * phi / pi
    out = ''
    # Check if the track hits the radial or longitudinal wall of the cavity
    if fabs(eta) < etatrans: # radial wall
        out += '\\psline[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]{->}(%10.5f;%10.5f)\n' % (rhomax/2.,phi_deg)
        out += '\\psline[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt](%10.5f;%10.5f)(%10.5f;%10.5f)\n' % (rhomax/2.,phi_deg,rhomax,phi_deg)
        out += '\\pscircle*[linecolor=black](%10.5f;%10.5f){1.5}\n' % (rhomax, phi_deg)
    else: # longitudinal wall; adjust rhomax accordingly
        rho = fabs(z_cut/sinh(eta))
        out += '\\psline[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]{->}(%10.5f;%10.5f)\n' % (rho/2.,phi_deg)
        out += '\\psline[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt](%10.5f;%10.5f)(%10.5f;%10.5f)\n' % (rho/2.,phi_deg,rho,phi_deg)
        out += '\\pscircle*[linecolor=black](%10.5f;%10.5f){1.5}\n' % (rho, phi_deg)
    return out

## Draws a straight track in rho-z space.
#
def drawstraighttrackrhoz(phi,eta,rhomax,zmax):
    etatrans = asinh(zmax/rhomax)
    phi_deg = 180. * phi / pi
    out = ''
    if fabs(eta) < etatrans: # radial wall hit.
        z = rhomax * sinh(eta)
        v = rhomax * sin(phi)
        #out += '\\psline[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]{->}(0.0,0.0)(%10.5f,%10.5f)\n' % (z, v)
        out += '\\psline[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]{->}(0.0,0.0)(%10.5f,%10.5f)\n' % (z/2., v/2.)
        out += '\\psline[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (z/2.,v/2., z, v)
        out += '\\pscircle*[linecolor=black](%10.5f,%10.5f){1.5}\n' % (z, v)
    else: # longitudinal wall
        z = (zmax * eta) / fabs(eta) # get the sign of the z right from the eta value
        v = (z / sinh(eta))*sin(phi)
        out += '\\psline[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]{->}(0.0,0.0)(%10.5f,%10.5f)\n' % (z/2., v/2.)
        out += '\\psline[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (z/2.,v/2.,z, v)
        out += '\\pscircle*[linecolor=black](%10.5f,%10.5f){1.5}\n' % (z, v)
    return out

## Draws the interaction point at the origin.
def drawinteractionpoint():
    out = ''
    out += '\\pscircle[linewidth=0.3pt,fillstyle=solid,fillcolor=white](0.0,0.0){1.0}\n'
    return out

## Draws a single tracking layer in rho-phi space.
#
def drawtrackerrhophi(rho,Deltarho,Nphi):
    out = ''
    if Deltarho < 0.00001:
        Deltarho = 0.01
    if Nphi == 0: # No segmentation in phi
        Deltarho = 1.5
        out += '\\pscircle[linewidth=0.3pt](%10.5f,%10.5f){%10.5f}\n' % (0.0,0.0,rho)
        out += '\\pscircle[linewidth=0.3pt](%10.5f,%10.5f){%10.5f}\n' % (0.0,0.0,rho+Deltarho)
    else: # Segmented in phi
        pass
        #for i in range(int(Nphi)):
        #    ## Azimuthal angle 1 in the ith segment.
        #    phi1 = i * (360. / Nphi)
        #    ## Azimuthal angle 2 in the ith segment.
        #    phi2 = phi1 + (360. / Nphi)
        #    # Draw the segment with the user thickness of minimum thickness if Delta_rho is too small.
        #    out += "\\pscustom[linewidth=0.3pt,linecolor=black]{\n"
        #    out += "\\psarc(%5.2f,%5.2f){%5.2f}{%5.2f}{%5.2f}\n" % (0.0,0.0,rho,phi1,phi2)
        #    out += "\\psarcn(%5.2f,%5.2f){%5.2f}{%5.2f}{%5.2f}\n" % (0.0,0.0,rho+0.01,phi2,phi1)
        #    out += "\\closepath\n"
        #    out += "}\n"
    return out

def drawmagfield(x,y,rad,magfield):
    out = ''
    out += '\\pscircle[linewidth=0.3pt](%10.5f,%10.5f){%10.5f}\n' % (x,y,rad)
    out += '\\uput{%10.5f}[-90.](%10.5f,%10.5f){\\scriptsize $B = %3.1f \\, \\textrm{T}$}\n' % (1.5*rad,x,y, magfield)
    out += '\\psline[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (x,y,x+rad/sqrt(2.),y+rad/sqrt(2.))
    out += '\\psline[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (x,y,x-rad/sqrt(2.),y+rad/sqrt(2.))
    out += '\\psline[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (x,y,x+rad/sqrt(2.),y-rad/sqrt(2.))
    out += '\\psline[linewidth=0.3pt](%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (x,y,x-rad/sqrt(2.),y-rad/sqrt(2.))
    return out

def drawmagfieldz(z,y,length,magfield):
    out = ''
    out += '\\psline[linewidth=0.3pt]{->}(%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (z-length,y,z+length,y)
    out += '\\uput{%10.5f}[-90.](%10.5f,%10.5f){\\scriptsize $B = %3.1f \\, \\textrm{T}$}\n' % (length/2.,z,y,magfield)
    return out


# Detector properties
#---------------------
z_det = 151.3 # cm
#Deltaz_det = 0.15
rho_det = 25 # cm
#Deltarho_det = 0.0
#Nphi_det = 32.
#Deltaphi_det = 2 * pi / Nphi_det
#Deltaphi_det_deg = 360./Nphi_det
#Eta_det_max = -1.0 * log( tan( (atan2(rho_det,z_det)/2.) ) )
Eta_det_max = asinh(z_det/rho_det)

# Cavity properties
#-------------------
z_cut = 282 # cm
rho_cut = 129 # cm
Eta_cut = asinh(z_cut/rho_cut)
print "Eta_cut = %5.3f" % Eta_cut

# Magnetic field
#----------------
B = 3.8 # Tesla

# Constants
#-----------
c = 299700000. # ms-1

# Figure properties
#-------------------
rhozfigwidth = 12.12 # cm
cutedge = 5.0 # cm
zcutplusedges = 2. * (z_cut + cutedge) # cm
scale = zcutplusedges / rhozfigwidth
#print "Scale for figures = %10.5f" % scale
rhocutplusedges = 2. * (rho_cut + cutedge) # cm
rhozfigheight = rhocutplusedges / scale
rhozyoffset = rhozfigheight * 0.5
rhozxoffset = rhozfigwidth * 0.5

rhophifigheight = rhozfigheight
rhophifigwidth  = rhozfigheight
rhophixoffset = rhophifigheight * 0.5
rhophiyoffset = rhophifigwidth * 0.5

rhophifilename = 'magrhophi'
rhozfilename = 'magrhoz'

## Current working directory.
cwd = os.getcwd()
## Temporary directory for files used in figure construction.
tempdir = tempfile.mkdtemp('.drawtracker')
## Path to the rhophi .tex file
rhophitexpath = os.path.join(tempdir, rhophifilename+'.tex')
## Path to the rhoz .tex file
rhoztexpath = os.path.join(tempdir, rhozfilename+'.tex')

## The Tex file to which the rho-phi figure is written.
rhophitexfile = open(rhophitexpath,'w')
## The Tex file to which the rho-z figure is written.
rhoztexfile = open(rhoztexpath,'w')

# Write the LaTeX document headers.
rhophitexfile.write(writeheader(rhophifigheight,rhophifigwidth,rhophixoffset,rhophiyoffset,scale))
rhoztexfile.write(writeheader(rhozfigheight,rhozfigwidth,rhozxoffset,rhozyoffset,scale))

# Draw cavity walls and the tracker
rhophitexfile.write(drawcutoffrhophi(rho_cut,cutedge))
rhophitexfile.write(drawtrackerrhophi(rho_det,0.0,0))
#
rhoztexfile.write(drawcutoffrhoz(z_cut,rho_cut,cutedge))
rhoztexfile.write(drawtrackerrhoz(z_det,0.0,rho_det,0.0))

# Axes
# x axis
rhophitexfile.write('\\psline[linewidth=0.3pt]{->}(0.0,0.0)(15.0,0.0)\n')
rhophitexfile.write('\\uput{4.0}[-90](15.0,0.0){\\scriptsize $x$}\n')
# y axis
rhoztexfile.write('\\psline[linewidth=0.3pt]{->}(0.0,0.0)(0.0,15.0)\n')
rhoztexfile.write('\\uput{2.0}[0](0.0,15.0){\\scriptsize $y$}\n')
# z axis
rhoztexfile.write('\\psline[linewidth=0.3pt]{->}(0.0,0.0)(15.0,0.0)\n')
rhoztexfile.write('\\uput{4.0}[-90](15.0,0.0){\\scriptsize $z$}\n')

# rho labels
rhocutlabelx = 35.0
gap = 15.
rhophitexfile.write('\\psline[linewidth=0.3pt]{|->|}(%10.5f,0.0)(%10.5f,%10.5f)\n' % (rhocutlabelx,rhocutlabelx,rho_cut))
rhophitexfile.write('\\uput{4.0}[0](%10.5f,%10.5f){\\scriptsize $\\rho \\,_{\\mathrm{cut}}$}\n' % (rhocutlabelx,rho_cut/2.))
rhophitexfile.write('\\psline[linewidth=0.3pt]{|->|}(%10.5f,0.0)(%10.5f,%10.5f)\n' % (rhocutlabelx+gap,rhocutlabelx+gap,rho_det))
rhophitexfile.write('\\uput{4.0}[0](%10.5f,%10.5f){\\scriptsize $\\rho \\,_{\\mathrm{det.}}$}\n' % (rhocutlabelx+gap,rho_det/2.))
rhocutlabelz = 200.0 # cm
rhoztexfile.write('\\psline[linewidth=0.3pt]{|->|}(%10.5f,0.0)(%10.5f,%10.5f)\n' % (rhocutlabelz,rhocutlabelz,rho_cut))
rhoztexfile.write('\\uput{4.0}[0](%10.5f,%10.5f){\\scriptsize $\\rho \\,_{\\mathrm{cut}}$}\n' % (rhocutlabelz,rho_cut/2.))
rhoztexfile.write('\\psline[linewidth=0.3pt]{|->|}(%10.5f,0.0)(%10.5f,%10.5f)\n' % (rhocutlabelz+gap,rhocutlabelz+gap,rho_det))
rhoztexfile.write('\\uput{4.0}[0](%10.5f,%10.5f){\\scriptsize $\\rho \\,_{\\mathrm{det.}}$}\n' % (rhocutlabelz+gap,rho_det/2.))

# z labels
zcutlabely = -114.0
rhoztexfile.write('\\psline[linewidth=0.3pt]{|->|}(0.0,%10.5f)(%10.5f,%10.5f)\n' % (zcutlabely,z_cut,zcutlabely))
rhoztexfile.write('\\uput{4.0}[-90](%10.5f,%10.5f){\\scriptsize $z \\,_{\\mathrm{cut}}$}\n' % (z_cut/2.,zcutlabely))
rhoztexfile.write('\\psline[linewidth=0.3pt]{|->|}(0.0,%10.5f)(%10.5f,%10.5f)\n' % (zcutlabely+gap,z_det,zcutlabely+gap))
rhoztexfile.write('\\uput{4.0}[-90](%10.5f,%10.5f){\\scriptsize $z \\,_{\\mathrm{det.}}$}\n' % (z_det/2.,zcutlabely+gap))

# Tracks
#phis = { 1 : pi/2., 2 : -pi/3., 3 : -5. * pi / 6. } # radians
phis = { 1 : pi/2., 2 : -pi/3., 3 : 2.76219239086  } # radians
etas = { 1 : 1.6193,   2 : -0.2,   3 : -2.6           }
ps   = { 1 : 0.200, 2 : .300,   3 : 0.161483595284 } # GeV
Ps   = {}
qs   = {}
Es   = { 0 : 0.0 }
Ms   = { 1 : .13957, 2 : .13957, 3 : .13957        } # GeV}
charges = { 0 : 1, 1 :  1, 2 : -1, 3 : 1 } # unit charge
pxs  = { 0 : 0.0 }
pys  = { 0 : 0.0 }
qs   = { 0 : 0.0 }
## Track radius of curvature (cm).
Rs   = {}
## Azimuthal angle relative to phi_trk that the particle hits the detector at rho_det (radians).
psidets = {}
## Time spent in the cavity
ts = {}
xhits = { 1 : [], 2 : [], 3 : [] }
yhits = { 1 : [], 2 : [], 3 : [] }
zhits = { 1 : [], 2 : [], 3 : [] }
for track, phi in phis.iteritems():
    #rhophitexfile.write(drawstraighttrackrhophi(phis[track],etas[track],rho_cut,z_cut))
    #rhoztexfile.write(drawstraighttrackrhoz(phis[track],etas[track],rho_cut,z_cut))
    pxs[track] = ps[track] * cos(phis[track])
    pys[track] = ps[track] * sin(phis[track])
    Ps[track]  = ps[track] * cosh(etas[track])
    qs[track]  = ps[track] * sinh(etas[track])
    Es[track]  = sqrt(Ps[track]*Ps[track] + Ms[track]*Ms[track])
    Es[0]  += Es[track]
    pxs[0] += pxs[track]
    pys[0] += pys[track]
    qs[0]  += qs[track]
    print "[P^{\\mu}] = (%10.5f, %10.5f, %10.5f, %10.5f)" % (Es[track], pxs[track], pys[track],  qs[track]  )
    print "[P^{\\mu}] = (%10.5f, %10.5f, %10.5f, %10.5f)" % (Ms[track], ps[track],  phis[track], etas[track])
    #
    Rs[track] = ps[track] / (0.003 * B * charges[track])
    print "R_trk = %10.5f" % (Rs[track])
    #
    # Does the track hit the radial wall?
    if (ps[track] > (0.003 * B * abs(charges[track]) * rho_cut)):
        print "Track %s hits the radial wall" % track
    else:
        print "Track %s misses the radial wall" % track
        # Calculate time spent before hitting the longitudinal wall
        ts[track] = (z_cut * Es[track]) / (c * fabs(qs[track]))
        print "Time particle %s spends in the tracker = %10.5f ns" % (track, ts[track]*10000000.)
        # Track intersecting with the tracking layer calculations
        # Does the track hit the tracking layer?
        if ( (2. * fabs(Rs[track])) > rho_det):
            print "Track %s hits the tracking layer at %10.5f" % (track, rho_det)
            psidets[track] = asin(rho_det / (2. * Rs[track]))
            print "Relative angle particle %s hits tracking layer = %10.5f" % (track, psidets[track])
            particlein = True; n = 0; phidets = []
            signpsi = psidets[track] / fabs(psidets[track])
            #
            circumx = -1.0 * Rs[track] * sin(phis[track])
            circumy =        Rs[track] * cos(phis[track])
            while (particlein):
                phidet = psidets[track] + (2. * pi * n * signpsi)
                tdet = (phidet * Rs[track] * Es[track]) / (ps[track] * c)
                xhit = rho_det * cos(phidet + phis[track])
                yhit = rho_det * sin(phidet + phis[track])
                rhohit = sqrt(xhit*xhit + yhit*yhit)
                #zhit = phidet * Rs[track] * sinh(etas[track])
                zhit = 2. * tdet * c * ((qs[track])/(Es[track]))
                if (tdet > ts[track]):
                    print "* BREAKING: tdet     = %10.5f ns - has already hit the cavity wall" % (tdet * 10000000.)
                    break
                if (fabs(zhit) > z_det):
                    print "* BREAKING: zhit     = %10.5f cm - track has left the tracking layer region" % (zhit)
                    break
                print "* phi_det        = %10.5f at t = %10.5f ns, x = %10.5f, y = %10.5f, z = %10.5f, rho    = %10.5f" % (phidet, tdet*10000000., xhit, yhit, zhit, rhohit)
                phidets.append(phidet)
                xhits[track].append(xhit)
                yhits[track].append(yhit)
                zhits[track].append(zhit)
                zhitplus = 2. * pi * fabs(Rs[track]) * sinh(etas[track])
                if fabs(zhit + zhitplus) < z_det:
                    phidets.append(phidet)
                    xhits[track].append(xhit)
                    yhits[track].append(yhit)
                    zhits[track].append(zhit + zhitplus)
                #
                # Complementary hit
                phidetcomp = pi * signpsi * ((2. * float(n)) + 1) - psidets[track]
                tdetcomp = (phidetcomp * Rs[track] * Es[track]) / (ps[track] * c)
                xhitcomp = rho_det * cos(phidetcomp + phis[track])
                yhitcomp = rho_det * sin(phidetcomp + phis[track])
                zhitcomp = 2. * phidetcomp * Rs[track] * sinh(etas[track])
                rhohitcomp = sqrt(xhitcomp*xhitcomp + yhitcomp*yhitcomp)
                if tdetcomp > ts[track]:
                    print "* BREAKING: tdetcomp = %10.5f ns - has already hit the cavity wall" % (tdet * 10000000.)
                    break
                if (fabs(zhitcomp) > z_det):
                    print "* BREAKING: zhitcomp = %10.5f cm - track has left the tracking layer region" % (zhitcomp)
                    break
                print "* phi_det (comp) = %10.5f at t = %10.5f ns, x = %10.5f, y = %10.5f, z = %10.5f, rhohit = %10.5f" % (phidetcomp, tdetcomp*10000000., xhitcomp, yhitcomp, zhitcomp, rhohitcomp)
                phidets.append(phidetcomp)
                xhits[track].append(xhitcomp)
                yhits[track].append(yhitcomp)
                zhits[track].append(zhitcomp)
                zhitcompplus = 2. * pi * fabs(Rs[track]) * sinh(etas[track])
                if fabs(zhitcomp + zhitcompplus) < z_det:
                    phidets.append(phidetcomp)
                    xhits[track].append(xhitcomp)
                    yhits[track].append(yhitcomp)
                    zhits[track].append(zhitcomp + zhitcompplus)
                n += 1
                #particlein = False
        else:
            print "Particle %s misses the tracking layer."
    print

#print "TOTALS"
#print "[P^{\\mu}] = (%10.5f, %10.5f, %10.5f, %10.5f)" % (Es[0], pxs[0], pys[0],  qs[0]  )
#ps[0] = sqrt(pxs[0]*pxs[0] + pys[0]*pys[0])
#Ps[0] = sqrt(ps[0]*ps[0] + qs[0]*qs[0])
#Ms[0] = sqrt(Es[0]*Es[0] - Ps[0]*Ps[0])
#print "[P^{\\mu}] = (%10.5f, %10.5f, %10.5f, %10.5f)" % (Ms[0], ps[0], 0.0, -10.0)
#print

# Draw the particle tracks
#--------------------------
its_per_s = 500000000
# Loop over the particles
for i, phi in phis.iteritems():
    if i == 0: continue
    print "TRACK %s" % i
    num_its = int(ts[i] * float(its_per_s)) + 1
    print "* Time spent in cavity = %10.5f ns (%s iterations)." % (ts[i] * 10000000., num_its)
    # rho-phi view
    rhophitexfile.write('\\pscurve[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]')
    rhoztexfile.write('\\pscurve[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]')
    # The circumcentre
    circumx = -1.0 * Rs[i] * sin(phis[i])
    circumy =        Rs[i] * cos(phis[i])
    angfreq = (ps[i] * c) / (Es[i] * Rs[i])
#                i = int(phi/Deltaphi_det)
#                j = int(z/(Deltaz_det))
#                maptexfile.write('\\psframe[linewidth=0.3pt,fillstyle=solid,fillcolor=lightgray]')
#                z_j   = float(j)*Deltaz_det
#                z_jp1 = z_j + ((z/fabs(z)) * Deltaz_det)
#                u_i   = float(i)*Deltaphi_det*rho_det
#                u_ip1 = u_i + ((phi/fabs(phi)) * Deltaphi_det * rho_det)
#                maptexfile.write('(%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (z_j,u_i,z_jp1,u_ip1))
#            rhoztexfile.write('(%10.5f,%10.5f)' % (z,y))
#            rhophitexfile.write('(%10.5f,%10.5f)' % (x,y))
#            if (rho > rho_det + 0.05):
#                break
    last_rho = -1.0
    for it in range(num_its):
        t_it = float(it) / float(its_per_s)
        #print ("%10.5f" % (t_it * 10000000.))
        angle = angfreq*t_it
        x = circumx + Rs[i] * cos(phis[i] - pi/2. + angle)
        y = circumy + Rs[i] * sin(phis[i] - pi/2. + angle)
        z = (qs[i] / Es[i]) * c * t_it
        rho = sqrt(x*x + y*y)
        phi = atan2(y,x)
        if (rho > rho_det and last_rho < rho_det) or (rho < rho_det and last_rho > rho_det):
            print "HIT at phi = %10.5f, z = %10.5f (rho = %10.5f)" % (phi, z, rho)
        if (fabs(angle) < (2.*pi)):
            rhophitexfile.write('(%10.5f,%10.5f)' % (x,y))
        rhoztexfile.write('(%10.5f,%10.5f)' % (z,y))
        last_rho = rho
    rhophitexfile.write('\n')
    rhoztexfile.write('\n')
    print
    print
    #
    # The wall points
    wallx = circumx + Rs[i] * cos(phis[i] - pi/2. + ((ps[i] * z_cut)/(Rs[i] * fabs(qs[i]))))
    wally = circumy + Rs[i] * sin(phis[i] - pi/2. + ((ps[i] * z_cut)/(Rs[i] * fabs(qs[i]))))
    signz = qs[i] / fabs(qs[i])
    rhophitexfile.write('\\pscircle[linewidth=0.3pt,fillstyle=solid,fillcolor=black](%10.5f,%10.5f){1.0}\n' % (wallx,wally))
    rhoztexfile.write('\\pscircle[linewidth=0.3pt,fillstyle=solid,fillcolor=black](%10.5f,%10.5f){1.0}\n' % (z_cut*signz,wally))


# The hits.
for i, yhit in yhits.iteritems():
    #if not i == 1: continue
    count = 0
    for y in yhit:
        rhophitexfile.write('\\pscircle[linewidth=0.3pt,fillstyle=solid,fillcolor=white]')
        rhophitexfile.write('(%10.5f,%10.5f){1.0}\n' % (xhits[i][count],y))
        rhoztexfile.write('\\pscircle[linewidth=0.3pt,fillstyle=solid,fillcolor=white]')
        rhoztexfile.write('(%10.5f,%10.5f){1.0}\n' % (zhits[i][count], y))
        count += 1

# Interaction Point
rhophitexfile.write(drawinteractionpoint())
rhoztexfile.write(drawinteractionpoint())

# Particle info
pinfo = ''
infx = -z_cut
infy = rho_cut
labelgap = 18.0
for track, M in Ms.iteritems():
    if not track==0:
        pinfo += '\\uput{%10.5f}[-45](%10.5f,%10.5f){\\psframebox*{\\scriptsize ' % (labelgap,infx,infy-(labelgap*float(track-1)))
        pinfo += '$[P^{\\mu}_{%s}] = (%+5.3f, \\, %+5.3f, \\, %+5.3f, \\, %+5.3f)$' % (track, Es[track], pxs[track], pys[track], qs[track])
        #pinfo += ', $M_{%s} = %5.3f$, $p_{%s} = %5.3f$' % (track, Ms[track], track, ps[track])
        pinfo += ', $p_{%s} = %5.3f$' % (track, ps[track])
        pinfo += '}}\n' 
    #else:
    #    pinfo += '\\uput{%10.5f}[-45](%10.5f,%10.5f){\\psframebox*{\\scriptsize ' % (labelgap,infx,infy-(3.*labelgap))
    #    pinfo += '$[P^{\\mu}_{%s}] = (%+5.3f, \\, %+5.3f, \\, %+5.3f, \\, %+5.3f)$' % ("\\mathrm{s}", Es[track], pxs[track], pys[track], qs[track])
    #    pinfo += '}}\n' 

rhoztexfile.write(pinfo)

# info - pseudo-rapidity and time
# Particle 1
etainfo1 = ''
etainfo1 += '\\uput{0.1}[-90](%10.5f,%10.5f){\\psframebox*{\\scriptsize ' % (z_det,-rho_det-10.)
etainfo1 += '$\\eta_{%s} = %+4.2f$, $t = %4.1f \\, \\textrm{ns}$' % (1, etas[1], ts[1] * 10000000.)
etainfo1 += '}}\n' 
rhoztexfile.write(etainfo1)
phiinfo1 = ''
phi1 = 100.
phiinfo1 += '\\uput{%10.5f}[%10.5f](%10.5f;%10.5f){\\scriptsize ' % (0.1, -90,2.*rho_cut/3.,phi1)
phiinfo1 += '$\\phi_{%s} = \\frac{\\pi}{2} \\, \\textrm{rad.}$}\n' % (1)
phiinfo1 += '\\uput{%10.5f}[%10.5f](%10.5f;%10.5f){\\scriptsize ' % (20., -90,2.*rho_cut/3.,phi1)
phiinfo1 += '$R_{%s} =  %4.1f\\, \\textrm{cm}$}\n' % (1, Rs[1])
rhophitexfile.write(phiinfo1)
# Particle 2
etainfo2 = ''
etainfo2 += '\\uput{0.1}[-90](%10.5f,%10.5f){\\psframebox*{\\scriptsize ' % (-z_det/8.,-rho_det-40.)
etainfo2 += '$\\eta_{%s} = %+4.2f$, $t = %4.1f \\, \\textrm{ns}$' % (2, etas[2], ts[2] * 10000000.)
etainfo2 += '}}\n' 
rhoztexfile.write(etainfo2)
phiinfo2 = ''
phi2 = -30.
phiinfo2 += '\\uput{%10.5f}[%10.5f](%10.5f;%10.5f){\\scriptsize ' % (0.1, -90,rho_cut/2.,phi2)
phiinfo2 += '$\\phi_{%s} = - \\frac{5 \\, \\pi}{6} \\, \\textrm{rad.}$}\n' % (2)
phiinfo2 += '\\uput{%10.5f}[%10.5f](%10.5f;%10.5f){\\scriptsize ' % (20., -90,rho_cut/2.,phi2)
phiinfo2 += '$R_{%s} =  %4.1f \\, \\textrm{cm}$}\n' % (2, Rs[2])
rhophitexfile.write(phiinfo2)
# Particle 3
etainfo3 = ''
etainfo3 += '\\uput{0.1}[-90](%10.5f,%10.5f){\\psframebox*{\\scriptsize ' % (-1.2 * z_det,-rho_det-20.)
etainfo3 += '$\\eta_{%s} = %+4.2f$, $t = %4.1f \\, \\textrm{ns}$' % (3, etas[3], ts[3] * 10000000.)
etainfo3 += '}}\n' 
rhoztexfile.write(etainfo3)
phiinfo3 = ''
phi3 = -140.
phiinfo3 += '\\uput{%10.5f}[%10.5f](%10.5f;%10.5f){\\scriptsize ' % (0.1, -90,3.*rho_cut/5.,phi3)
phiinfo3 += '$\\phi_{%s} = %4.2f \\, \\textrm{rad.}$}\n' % (3, phis[3])
phiinfo3 += '\\uput{%10.5f}[%10.5f](%10.5f;%10.5f){\\scriptsize ' % (20., -90,3.*rho_cut/5.,phi3)
phiinfo3 += '$R_{%s} =  %4.1f \\, \\textrm{cm}$}\n' % (3, Rs[3])
rhophitexfile.write(phiinfo3)
#phiinfo2 = ''
#phiinfo2 += '\\uput{%10.5f}[%10.5f](%10.5f;%10.5f){\\scriptsize ' % (5.0, (180.*phis[2]/pi)+90,rho_cut/3.,180.*phis[2]/pi)
#phiinfo2 += '$\\phi_{%s} = -\\frac{5 \\, \\pi}{6} \\, \\textrm{rad.}$}\n' % (2)
#rhophitexfile.write(phiinfo2)
## particle 3
#phiinfo3 = ''
#phiinfo3 += '\\uput{%10.5f}[%10.5f](%10.5f;%10.5f){\\scriptsize ' % (5.0, (180.*phis[3]/pi)-90,2.*rho_cut/3.,180.*phis[3]/pi)
#phiinfo3 += '$\\phi_{%s} = %4.2f \\, \\textrm{rad.}$}\n' % (3, phis[3])
#rhophitexfile.write(phiinfo3)

# Magnetic field labels
rhophitexfile.write(drawmagfield(0.0,-100.,5.,B))
rhoztexfile.write(drawmagfieldz(-200.,-100.,-20.,B))


####################################################################
## DRAW THE PARTICLES
####################################################################
#from ROOT import TFile, TTree

#myfile = TFile("ntuple.root","READ")

##myfile.ls()

#mydir = myfile.Get('tomTree')
##mydir.ls()
#mychain = mydir.Get('tree')
#entries = mychain.GetEntriesFast()
##print entries

##branch = mychain.GetBranch("Process")
##myprocess = branch.GetAddress()
##print branch

##for jentry in xrange(entries):
#for jentry in xrange(1):
#    #print jentry
#    ientry = mychain.LoadTree(jentry)
#    if ientry < 0:
#        break
#    nb = mychain.GetEntry(jentry)
#    #print "nb = ", nb
#    if nb <= 0:
#        continue
#    #print
#    #print "Size of Es vector = %s" % mychain.Es.size()
#    #print "Size of Handles vector = %s" % mychain.Handles.size()
#    #print "Size of FinalState vector = %s" % mychain.Visible.size()
#    #print "Size of Visible vector = %s" % mychain.Visible.size()
#    #print mychain.Process
#    #print mychain.Generator
#    
#    #handles = {} # Dictionary for the handles
#    Energies = {}
#    pxs = {}
#    pys = {}
#    qs = {}
#    charges = {}
#    ids = {}
#    counter = 0
#    SE = 0.
#    Spx = 0.
#    Spy = 0.
#    Sq = 0.
#    ## Loop through the handles for the collection of interest
#    for h in mychain.FinalState:
#        #print "i = %4d, handle = %4d" % (counter,h)
#        counter += 1
#        Energies[h] = mychain.Es[counter]
#        pxs[h] = mychain.pxs[counter]
#        pys[h] = mychain.pys[counter]
#        qs[h]  = mychain.pzs[counter]
#        charges[h] = mychain.Charges[counter]
#        ids[h] = mychain.IDs[counter]
#        SE += mychain.Es[counter]
#        Spx += mychain.pxs[counter]
#        Spy += mychain.pys[counter]
#        Sq  += mychain.pzs[counter]
#    print "Total 4-Momentum = (%10.5f, %10.5f, %10.5f, %10.5f)" % (SE,Spx,Spy,Sq)

#    # Get the first visible particle
#    for handle in mychain.Visible:
#        #handle = mychain.Visible[1]
#        E = Energies[handle]
#        px = pxs[handle]
#        py = pys[handle]
#        q = qs[handle]
#        charge = charges[handle]
#        if (charge==0):
#            continue
#        ID = ids[handle]
#        p_trk = sqrt(px*px + py*py)
#        P = sqrt(px*px + py*py + q*q)
#        M = sqrt(E*E - P*P)
#        eta_trk = 0.5  * log((P + q)/(P - q))
#        if (p_trk < 0.15) or (fabs(eta_trk) > 2.5):
#            continue
#        theta_trk = atan2(p_trk,q)
#        theta_trk_deg = 180. * theta_trk / pi
#        phi_trk = atan2(py, px)
#        phi_trk_deg = 180. * phi_trk / pi
#        Beta = P / (sqrt(M*M + P*P))
#        Beta_t = Beta * sin(theta_trk)
#        Beta_z = Beta * cos(theta_trk)
#        c = 300000000. # speed of light
#        t_trk = z_det / (fabs(Beta_z) * c)
#        its_per_s = 10000000000
#        num_its = int(t_trk * float(its_per_s))
#        B = 3.8 # magnetic field [Tesla]
#        R_trk = (p_trk) / (0.3 * charge * B)
#        rad_travelled = ((Beta_t * c) / (fabs(R_trk))) * t_trk
#        angfreq = (Beta_t * c) / R_trk
#        #rhophitexfile.write('\\psline[linewidth=0.3pt]{->}(0.0,0.0)(%10.5f;%10.5f)\n' % (fabs(R_trk),phi_trk_deg))
#        ##R = 1./C
#        # The circumcentre
#        circumx = -1.0 * R_trk * sin(phi_trk)
#        circumy =        R_trk * cos(phi_trk)
#        rhoztexfile.write('\\pscurve[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]')
#        rhophitexfile.write('\\pscurve[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]')
#        last_rho = -0.1
#        for k in range(num_its):
#            t = k / float(its_per_s)
#            x = circumx + R_trk * cos(phi_trk - pi/2. + angfreq*t)
#            y = circumy + R_trk * sin(phi_trk - pi/2. + angfreq*t)
#            z = Beta_z * c * t
#            rho = sqrt(x*x + y*y)
#            phi = atan2(y,x)
#            if (rho > rho_det and last_rho < rho_det) or (rho < rho_det and last_rho > rho_det):
#                print "HIT at phi = %10.5f, z = %10.5f" % (phi, z)
#                i = int(phi/Deltaphi_det)
#                j = int(z/(Deltaz_det))
#                maptexfile.write('\\psframe[linewidth=0.3pt,fillstyle=solid,fillcolor=lightgray]')
#                z_j   = float(j)*Deltaz_det
#                z_jp1 = z_j + ((z/fabs(z)) * Deltaz_det)
#                u_i   = float(i)*Deltaphi_det*rho_det
#                u_ip1 = u_i + ((phi/fabs(phi)) * Deltaphi_det * rho_det)
#                maptexfile.write('(%10.5f,%10.5f)(%10.5f,%10.5f)\n' % (z_j,u_i,z_jp1,u_ip1))
#            rhoztexfile.write('(%10.5f,%10.5f)' % (z,y))
#            rhophitexfile.write('(%10.5f,%10.5f)' % (x,y))
#            if (rho > rho_det + 0.05):
#                break
#            last_rho = rho
#        rhoztexfile.write('\n')
#        rhophitexfile.write('\n')
#        ##addedphi_deg = 2.*(180./pi)*(asin(cutoff/(2.*R)))
#        ##rhophitexfile.write('\\pscircle[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt](%10.5f,%10.5f){%10.5f}\n' % (circumx,circumy,R_trk))
#        #rhophitexfile.write('\\psarc[linewidth=0.3pt,linestyle=dashed,dash=1.5pt 0.75pt]')
#        #if charge > 0:
#        #    rhophitexfile.write('(%10.5f,%10.5f){%10.5f}{%10.5f}{%10.5f}\n' % (circumx,circumy,R_trk,phi_trk_deg-90.,phi_trk_deg-90.+(rad_travelled*180./pi)))
#        #if charge < 0:
#        #    rhophitexfile.write('(%10.5f,%10.5f){%10.5f}{%10.5f}{%10.5f}\n' % (circumx,circumy,R_trk,phi_trk_deg-90.-(rad_travelled*180./pi),phi_trk_deg-90.))
#        ## Loop over z values
#        #rhoztexfile.write('\\pscurve[linewidth=0.3pt,linestyle=dashed,dash = 1.5pt 0.75pt]')
#        #for k in range(101):
#        #    z = z_det*(k/100.)
#        #    z = copysign(z,eta_trk)
#        #    psi = z * (tan(theta_trk) / (R_trk))       
#        #    y = circumy + R_trk*sin(phi_trk - pi/2. + psi)
#        #    #print z
#        #    rhoztexfile.write('(%10.5f,%10.5f)' % (z,y))
#        #rhoztexfile.write('\n')
#        #rhoztexfile.write('\\psline[linewidth=0.3pt]{->}(0.0,0.0)(%10.5f;%10.5f)\n' % (fabs(q),theta_trk_deg))
#    
####################################################################
#    

# Write the LaTeX document footer.
rhophitexfile.write(writefooter())
rhoztexfile.write(writefooter())
#maptexfile.write(writefooter())
# Close the LaTeX file.
rhophitexfile.close()
rhoztexfile.close()
#maptexfile.close()

#Make the PDF from the latex.
import subprocess

#run LaTeX on the generated figure LaTeX file.
#print texpath
#texcmd = ["latex", texpath]

#filenames = [rhophifilename, rhozfilename, mapfilename]
filenames = [rhophifilename, rhozfilename]
for filename in filenames:
    ## Process for running the LaTeX figure generation.
    texproc = subprocess.Popen(["latex", "-interaction=batchmode", os.path.join(tempdir,filename+".tex")], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=tempdir)
    texproc.wait()
    
    ## Process for running the dvi to ps conversion.
    dvproc = subprocess.Popen(["dvips", "-q", os.path.join(tempdir,filename+".dvi")], stdout=subprocess.PIPE, cwd=tempdir)
    dvproc.wait()
    
    ## Process for running the ps to pdf conversion.
    cnvproc = subprocess.Popen(["ps2pdf", os.path.join(tempdir,filename+".ps")], stdout=subprocess.PIPE, cwd=tempdir)
    cnvproc.wait()
    
    ## The output path for the final PDF.
    outpath = os.path.join(tempdir, filename+".pdf")
    if os.path.exists(outpath):
        shutil.copy(outpath, cwd)
    ## The output path for the final .tex file (for debugging).
    outtex = os.path.join(tempdir, filename+".tex")
    if os.path.exists(outtex):
        shutil.copy(outtex, cwd)

shutil.rmtree(tempdir, ignore_errors=True) # delete the temporary files

##print "Eta_det_max = ", Eta_det_max

#print "Visible particle Handle is ", handle
#print "[P^\\mu] = (E, px, py, q)   = (%10.5f, %10.5f, %10.5f, %10.5f)\n" % (E,px,py,q)
#print "         = (M, p, phi, eta) = (%10.5f, %10.5f, %10.5f, %10.5f)\n" % (M, p_trk, phi_trk, eta_trk)
#print "Charge = ", charge
#print "ID = ",ID
#print "R_trk = ", R_trk
#print "eta_trk = %10.5f, theta_trk = %10.5f (%10.5f deg)" %(eta_trk, theta_trk, theta_trk_deg)
#print "Beta, Beta_t, Beta_z = ", Beta, Beta_t, Beta_z
#print "t_trk = %10.7f ns" % (t_trk*1000000000.)
#print "number of iterations = %s" % num_its
##print "radians travelled = ", rad_travelled
